from data.dataset import DATASET_REGISTRY
import time
import argparse

def parse_args():
    parser = argparse.ArgumentParser(description="prepare dataset")
    parser.add_argument("dataset", type=str, help='dataset name')
    args = parser.parse_args()
    return args

def prepare_dataset(dataset_name):
    print(">>>> {} <<<<".format(dataset_name))
    start_time = time.time()
    dataset_dicts = DATASET_REGISTRY.get(dataset_name)()
    cost_time = time.time() - start_time
    print("cost time: {} s".format(cost_time))
    print("len: {}".format(len(dataset_dicts)))

if __name__ == '__main__':
    args = parse_args()
    prepare_dataset(args.dataset + '-train')
    prepare_dataset(args.dataset + '-test')
    print("done")